#pragma once
#include <vector>
#include <map>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <iostream>
#include <initializer_list>
using std::vector;
using std::map;
using std::string;
using std::unordered_map;
using std::unordered_set;
using std::initializer_list;

//并查集的数据结构实现
namespace DSU
{
	template<class T>
	class UnionFindSet
	{
	public:
		// 默认构造函数
		UnionFindSet() = default;

		// 列表初始化构造函数
		UnionFindSet(std::initializer_list<T> initList)
		{
			size_t i = 0;
			for (const T& element : initList)
			{
				_data.push_back(element);
				_indexMap[element] = i;
				_parent.push_back(i);
				_size.push_back(1);
				++i;
			}
			_setCount = initList.size();
		}

		// 通过数组初始化
		UnionFindSet(const T* a, size_t n)
		{
			for (size_t i = 0; i < n; ++i)
			{
				_data.push_back(a[i]);
				_indexMap[a[i]] = i;
				_parent.push_back(i);
				_size.push_back(1);
			}
			_setCount = n;
		}

		// 通过vector初始化
		UnionFindSet(const vector<T>& data)
		{
			for (size_t i = 0; i < data.size(); ++i)
			{
				_data.push_back(data[i]);
				_indexMap[data[i]] = i;
				_parent.push_back(i);
				_size.push_back(1);
			}
			_setCount = data.size();
		}

		// 添加新元素
		void AddElement(const T& element)
		{
			if (_indexMap.find(element) != _indexMap.end())
				return;

			size_t newIndex = _data.size();
			_data.push_back(element);
			_indexMap[element] = newIndex;
			_parent.push_back(newIndex);
			_size.push_back(1);
			_setCount++;
		}

		// 批量添加元素
		void AddElements(std::initializer_list<T> elements)
		{
			for (const T& element : elements)
			{
				AddElement(element);
			}
		}

		// 查找元素所在集合的根节点索引
		int FindRoot(int index)
		{
			// 路径压缩
			if (_parent[index] != index)
			{
				_parent[index] = FindRoot(_parent[index]);
			}
			return _parent[index];
		}

		// 查找元素所在集合的根节点索引（通过元素值）
		int FindRoot(const T& element)
		{
			auto it = _indexMap.find(element);
			if (it == _indexMap.end())
			{
				return -1;
			}
			return FindRoot(it->second);
		}

		// 合并两个元素所在的集合
		void Union(const T& a, const T& b)
		{
			int rootA = FindRoot(a);
			int rootB = FindRoot(b);

			if (rootA == -1 || rootB == -1 || rootA == rootB)
				return;

			// 按秩合并：将小树合并到大树上
			if (_size[rootA] < _size[rootB])
			{
				_parent[rootA] = rootB;
				_size[rootB] += _size[rootA];
			}
			else
			{
				_parent[rootB] = rootA;
				_size[rootA] += _size[rootB];
			}
			_setCount--;
		}

		// 批量合并
		void Union(std::initializer_list<std::pair<T, T>> pairs)
		{
			for (const auto& pair : pairs)
			{
				Union(pair.first, pair.second);
			}
		}

		// 判断两个元素是否属于同一集合
		bool IsSameSet(const T& a, const T& b)
		{
			int rootA = FindRoot(a);
			int rootB = FindRoot(b);
			return (rootA != -1 && rootB != -1 && rootA == rootB);
		}

		// 获取集合个数
		size_t GetSetCount() const
		{
			return _setCount;
		}

		// 根据集合根节点显示集合内所有元素
		void DisplaySetByRoot(const T& rootElement)
		{
			int rootIndex = FindRoot(rootElement);
			if (rootIndex == -1)
			{
				std::cout << "元素不存在或不是根节点!" << std::endl;
				return;
			}

			std::cout << "集合根节点 '" << _data[rootIndex] << "' 包含的元素:" << std::endl;
			for (size_t i = 0; i < _data.size(); ++i)
			{
				if (FindRoot(i) == rootIndex)
				{
					std::cout << "  编号 " << i << ": " << _data[i] << std::endl;
				}
			}
		}

		// 显示所有集合
		void DisplayAllSets()
		{
			std::cout << "总共 " << _setCount << " 个集合:" << std::endl;

			unordered_set<int> displayedRoots;

			for (size_t i = 0; i < _data.size(); ++i)
			{
				int root = FindRoot(i);
				if (displayedRoots.find(root) == displayedRoots.end())
				{
					displayedRoots.insert(root);
					std::cout << "集合 " << root << " (根节点: '" << _data[root] << "'): ";

					vector<T> elements;
					for (size_t j = 0; j < _data.size(); ++j)
					{
						if (FindRoot(j) == root)
						{
							elements.push_back(_data[j]);
						}
					}

					for (size_t k = 0; k < elements.size(); ++k)
					{
						std::cout << elements[k];
						if (k < elements.size() - 1) std::cout << ", ";
					}
					std::cout << std::endl;
				}
			}
		}

		// 根据编号查找元素所在集合的根节点内容
		T GetSetRootByIndex(int index)
		{
			if (index < 0 || index >= _data.size())
			{
				return T();
			}
			int rootIndex = FindRoot(index);
			return _data[rootIndex];
		}

		// 根据元素值查找所在集合的根节点内容
		T GetSetRootByElement(const T& element)
		{
			int rootIndex = FindRoot(element);
			if (rootIndex == -1) return T();
			return _data[rootIndex];
		}

		// 获取集合大小
		size_t GetSetSize(const T& element)
		{
			int rootIndex = FindRoot(element);
			if (rootIndex == -1) return 0;
			return _size[rootIndex];
		}

		// 获取所有元素
		const vector<T>& GetElements() const
		{
			return _data;
		}

		// 检查元素是否存在
		bool Contains(const T& element) const
		{
			return _indexMap.find(element) != _indexMap.end();
		}

	private:
		vector<T> _data;
		map<T, int> _indexMap;
		vector<int> _parent;
		vector<int> _size;
		size_t _setCount = 0;
	};
}
